from neo4j.graph import Node, Relationship, Path
from neo4j.time import DateTime

"""
class for parsing neo4j graph data into frontend format
"""

color_menu = {
    "red": "#FF0000",
    "orange": "#FE9900",
    # "orange": "#d89c7a",  # 莫兰迪橙色
    "grey": "#808080",
    "black": "#000000",
    "white": "#FFFFFF",
    "purple": "#800080",
}

# label颜色，数据越少的颜色越浅
label_color_menu = {
    "Phone": "#306F88",     # 3750
    "Device": "#8cacc5",     # 3750
    "User": " #90795D",     # 3000
    "Addr": "#94B267",      # 1186
    "LAC": "#D5C095",       # 215
    "APP": "#DEDDCB",       # 36
    "Product": "#D9EAD7",   # 10
    "Group": "#F8DACC",     # 10
    "BrdAddr": "#94B267",
    "News": "#e8a3c1",
    "Industry": "#d294bc",
    "Keyword": "#ffbc5e",
}

# mapping neo4j node label to business id key name in properties
id_map = {
    "User": "user_id",
    "Phone": "mobile_number",
    "Device": "imei",
    "LAC": "lac_code",
    "APP": "app_name",
    "Product": "product_name",
    "Group": "group_name",
    "Addr": "addr",
    "BrdAddr": "addr",
    "News": "news_title",
    "Industry": "industry_name",
    "Keyword": "keyword",
}


def _get_node_buzid(node: Node):
    node_label = list(node.labels)[0]
    return (node_label, str(dict(node.items())[id_map[node_label]]).translate(str.maketrans('', '', '@#$%^&*()_+,.')).
            replace(' ', ''))


def add_vertex(node: Node, vertex_set, key):
    node_label, node_buz_id = _get_node_buzid(node)
    properties = dict(node.items())
    if node_buz_id not in vertex_set:
        vertex_set[node_buz_id] = {
            "id": node_buz_id,
            "label": node_label,
            "properties": properties,
        }
        if "group_fam_id" in properties:
            properties["groupId"] = properties["group_fam_id"]
        elif "group_geo_id" in properties:
            properties["groupId"] = properties["group_geo_id"]
        _convert_neo4j_date(properties)
    _add_color(key, vertex_set[node_buz_id])


def add_edge(edge: Relationship, edge_set, key):
    if not edge.start_node or not edge.end_node:
        return
    start_node_label, start_node_buz_id = _get_node_buzid(edge.start_node)
    end_node_label, end_node_buz_id = _get_node_buzid(edge.end_node)
    edge_buz_id = f"{start_node_buz_id}-{edge.type}-{end_node_buz_id}-{str(edge.element_id)[-4:]}".translate(
        str.maketrans('', '', '@#$%^&*()_+:'))
    if edge_buz_id not in edge_set:
        edge_set[edge_buz_id] = {
            "id": edge_buz_id,
            "srcId": start_node_buz_id,
            "dstId": end_node_buz_id,
            "label": edge.type,
            "properties": dict(edge.items()),
        }
        _convert_neo4j_date(edge_set[edge_buz_id]['properties'])
    _add_color(key, edge_set[edge_buz_id])


def _add_color(key, entity):
    """为点或边添加前端的强调色"""
    # 添加基础颜色
    if entity["label"] in label_color_menu and "color" not in entity:
        entity["color"] = label_color_menu[entity["label"]]

    # 添加强调颜色
    if "_" in key and key[:key.index("_")] in color_menu:
        color_name = key[:key.index("_")]
        entity["color"] = color_menu[color_name]




def _convert_neo4j_date(properties: dict):
    """转化一个dict中类型为neo4j.time.DateTime的值为string"""
    for k in properties:
        if isinstance(properties[k], DateTime):
            properties[k] = str(properties[k])[:20]
    if 'weight' in properties:
        del properties['weight']


def evaluate_value(value, vertex_set, edge_set, other, key=None):
    if isinstance(value, Node):
        # print(f"Node {key} recognised: {value}")
        add_vertex(value, vertex_set, key)
    elif isinstance(value, Relationship):
        # print(f"Relationship  {key}  recognised: {value}")
        add_edge(value, edge_set, key)
    elif isinstance(value, Path):
        # print(f"Path  {key}  recognised: {value}")
        for node in value.nodes:
            # print(f"Node recognised: {node}")
            add_vertex(node, vertex_set, key)
        for rel in value.relationships:
            # print(f"Relationship recognised: {rel}")
            add_edge(rel, edge_set, key)
    elif isinstance(value, list):
        # print(f"list  {key}  recognised: {value}")
        for elem in value:
            evaluate_value(elem, vertex_set, edge_set, other, key)

    # 推理边逻辑
    elif isinstance(value, dict) and "prediction" in key and "start_node" in value and "end_node" in value:
        start_node: Node = value["start_node"]
        end_node: Node = value["end_node"]
        start_node_label, start_node_buz_id = _get_node_buzid(start_node)
        end_node_label, end_node_buz_id = _get_node_buzid(end_node)
        edge_buz_id = f"{start_node_buz_id}-{end_node_buz_id}"
        edge_set[edge_buz_id] = {
            "id": edge_buz_id,
            "srcId": start_node_buz_id,
            "dstId": end_node_buz_id,
            "label": "prediction",
            "properties": dict(),
            "dashed": 1
        }
    else:  # 其他类型的字段
        if key not in other:
            other[key] = set()
        other[key].add(str(value))


def parse_graph(records, keys):
    vertex_set, edge_set, other = {}, {}, {}
    for record in records:
        for i in range(len(record)):
            key = keys[i]
            value = record[i]
            evaluate_value(value, vertex_set, edge_set, other, key)
    # print(f"graph parsed vcnt, ecnt, othercnt: {len(vertex_set), len(edge_set), len(other)}")

    return vertex_set, edge_set, other


def simplify_for_explain(vertex_set, edge_set):
    """
    为大模型解释进一步简化图结构，返回最简点边列表
    :param vertex_set:
    :param edge_set:
    :return:
    """
    vertex_list = []
    edge_list = []
    for vertexId in vertex_set:
        vertex_for_explain = vertex_set[vertexId].copy()
        vertex_list.append(vertex_for_explain)

    for edgeId in edge_set:
        edge_for_explain = edge_set[edgeId].copy()
        del edge_for_explain['id']
        edge_list.append(edge_for_explain)

    return vertex_list, edge_list
